package rsa

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"github.com/gogf/gf/errors/gerror"
	"github.com/gogf/gf/util/gconv"
)

// Keypair RSA公私钥对类型
type Keypair struct {
	PrivateKey string `json:"privateKey"` // base64编码后的私钥
	PublicKey  string `json:"publicKey"`  // base64编码后的公钥
}

// NewKeypair 生成RSA公私钥, 公私钥均使用base64编码后返回
func NewKeypair() (*Keypair, error) {
	// 生成密钥对
	keypairPrivateKey, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return nil, err
	}
	// 生成私钥
	privateKey, err := x509.MarshalPKCS8PrivateKey(keypairPrivateKey)
	if err != nil {
		return nil, err
	}
	// 生成公钥, 注意这里传递指针
	publicKey, err := x509.MarshalPKIXPublicKey(&keypairPrivateKey.PublicKey)
	if err != nil {
		return nil, err
	}
	return &Keypair{
		PrivateKey: base64.StdEncoding.EncodeToString(privateKey),
		PublicKey:  base64.StdEncoding.EncodeToString(publicKey),
	}, nil
}

// EncryptEverything 将传入数据转为json格式后进行加密, 返回base64编码的密文
func EncryptEverything(data interface{}, publicKey string) (string, error) {
	// 将数据转为json格式
	d, err := json.Marshal(data)
	if err != nil {
		return "", err
	}
	// 对json格式的数据进行加密
	encrypt, err := Encrypt(d, publicKey)
	if err != nil {
		return "", err
	}
	// 返回base64编码的密文
	return base64.StdEncoding.EncodeToString(encrypt), nil
}

// Encrypt 使用base64编码的公钥对数据进行加密
func Encrypt(data []byte, publicKey string) ([]byte, error) {
	// 获取公钥
	pub, err := parsePublicKey(publicKey)
	if err != nil {
		return nil, err
	}
	// 对数据进行加密
	return rsa.EncryptPKCS1v15(rand.Reader, pub, data)
}

// DecryptEverything 将密文通过base64解码后进行解密, 内部使用gconv.Scan()函数将数据转换并赋值给pointer
func DecryptEverything(encrypt string, privateKey string, pointer interface{}) error {
	// 将密文通过base64解码
	e, err := base64.StdEncoding.DecodeString(encrypt)
	if err != nil {
		return err
	}
	// 对数据进行解密
	decrypt, err := Decrypt(e, privateKey)
	if err != nil {
		return err
	}
	// 将转换并赋值给pointer
	return gconv.Scan(decrypt, pointer)
}

// Decrypt 使用base64编码的私钥对公钥加密后的数据进行解密
func Decrypt(encrypt []byte, privateKey string) ([]byte, error) {
	// 获取私钥
	pri, err := parsePrivateKey(privateKey)
	if err != nil {
		return nil, err
	}
	// 对数据进行解密
	return rsa.DecryptPKCS1v15(rand.Reader, pri, encrypt)
}

// SignEverything 将传入数据转为json格式后进行签名, 返回base64编码的签名
func SignEverything(data interface{}, privateKey string) (string, error) {
	// 将数据转为json格式
	d, err := json.Marshal(data)
	if err != nil {
		return "", err
	}
	// 对json格式的数据进行签名
	signature, err := Sign(d, privateKey)
	if err != nil {
		return "", err
	}
	// 返回base64编码的签名
	return base64.StdEncoding.EncodeToString(signature), nil
}

// Sign 通过base64编码后的私钥对数据进行签名, 使用sha256算法
func Sign(data []byte, privateKey string) ([]byte, error) {
	// 计算数据的sha256值
	h := sha256.New()
	h.Write(data)
	hashed := h.Sum(nil)
	// 获取私钥
	pri, err := parsePrivateKey(privateKey)
	if err != nil {
		return nil, err
	}
	// 对数据进行签名
	return rsa.SignPKCS1v15(rand.Reader, pri, crypto.SHA256, hashed)
}

// VerifyEverything 将传入数据转为json格式并将签名通过base64解码后进行验签
func VerifyEverything(data interface{}, signature string, publicKey string) (bool, error) {
	// 将传入数据转为json格式
	d, err := json.Marshal(data)
	if err != nil {
		return false, err
	}
	// 将签名通过base64解码
	s, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return false, err
	}
	return Verify(d, s, publicKey)
}

// Verify 通过base64编码后的公钥对时使用sha256算法签名的数据进行验签
func Verify(data []byte, signature []byte, publicKey string) (bool, error) {
	// 获取公钥
	pub, err := parsePublicKey(publicKey)
	if err != nil {
		return false, err
	}
	// 获取原始数据的sha256值
	hashed := sha256.Sum256(data)
	// 对数据进行验签
	err = rsa.VerifyPKCS1v15(pub, crypto.SHA256, hashed[:], signature)
	if err != nil {
		return false, err
	}
	return true, nil
}

// parsePrivateKey 将base64编码的私钥转为(*rsa.PrivateKey)类型的私钥
func parsePrivateKey(privateKey string) (*rsa.PrivateKey, error) {
	// 将base64编码的私钥解码
	pkb, err := base64.StdEncoding.DecodeString(privateKey)
	if err != nil {
		return nil, err
	}
	// 通过privateKey生成对应的PrivateKey对象
	pk, err := x509.ParsePKCS8PrivateKey(pkb)
	if err != nil {
		return nil, err
	}
	pri, ok := pk.(*rsa.PrivateKey)
	if !ok {
		return nil, gerror.Newf("私钥错误, 私钥的类型为(%T)", pk)
	}
	return pri, nil
}

// parsePublicKey 将base64编码的公钥转为(*rsa.PublicKey)类型的公钥
func parsePublicKey(publicKey string) (*rsa.PublicKey, error) {
	// 将base64编码的公钥解码
	pkb, err := base64.StdEncoding.DecodeString(publicKey)
	if err != nil {
		return nil, err
	}
	// 通过publicKey生成对应的PublicKey对象
	pk, err := x509.ParsePKIXPublicKey(pkb)
	if err != nil {
		return nil, err
	}
	pub, ok := pk.(*rsa.PublicKey)
	if !ok {
		return nil, gerror.Newf("公钥错误, 公钥的类型为(%T)", pk)
	}
	return pub, nil
}
